"""
Dify API服务
"""
import uuid
import requests
from config import DIFY_CONFIG
from loguru import logger

class DifyService:
    """Dify API服务类"""
    
    def __init__(self):
        """初始化Dify API服务"""
        self.base_url = DIFY_CONFIG['base_url']
        self.api_key = DIFY_CONFIG['api_key']
        self.endpoint = f"{self.base_url}{DIFY_CONFIG['endpoint']}"
        self.user = DIFY_CONFIG['user']
        self.response_mode = DIFY_CONFIG['response_mode']
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
    
    def get_data_sources(self, query_text):
        """
        通过Dify API获取数据源信息
        
        Args:
            query_text: 用户查询文本
        
        Returns:
            dict/None: API返回的解析后数据，失败则返回None
        """
        try:
            # 构造请求的提示语
            prompt = f"以下是一个搜索查询，请帮我找出相关的数据源，以便进行爬取：{query_text}\n\n请以JSON格式返回以下信息:\n1. 可能的入口URL(entry_url)\n2. 数据类别(category)，如法规与政策、新闻等\n3. 信源描述(description)\n4. 优先级(priority)：高、中、低\n5. 信源类型(source_type)：网页、API等\n6. 检索条件(retrieval_conditions)：关键词列表、过滤条件、日期范围等\n7. 爬取策略(scraping_strategies)：爬取频率、要提取的字段等\n8. 授权信息(authorization_info)：如果需要的话"
            
            # 构造请求负载
            payload = {
                "inputs": {},
                "query": prompt,
                "response_mode": self.response_mode,
                "conversation_id": str(uuid.uuid4()),
                "user": self.user
            }
            
            logger.info(f"发送请求到Dify API，查询文本: {query_text}")
            # 发送请求
            response = requests.post(self.endpoint, headers=self.headers, json=payload, timeout=60)
            
            # 检查响应状态
            if response.status_code == 200:
                result = response.json()
                logger.info("成功获取Dify API响应")
                
                # 从响应中提取返回的数据
                if 'answer' in result:
                    return self._parse_dify_response(result['answer'], query_text)
                else:
                    logger.warning("Dify API响应中没有answer字段")
                    return None
            else:
                logger.error(f"Dify API请求失败: 状态码 {response.status_code}, 响应 {response.text}")
                return None
        except requests.exceptions.RequestException as e:
            logger.error(f"Dify API请求异常: {e}")
            return None
        except Exception as e:
            logger.error(f"处理Dify API响应时发生错误: {e}")
            return None
    
    def _parse_dify_response(self, response_text, query_text):
        """
        解析Dify API返回的文本响应为结构化数据
        
        Args:
            response_text: Dify API返回的响应文本
            query_text: 原始查询文本
        
        Returns:
            dict: 解析后的结构化数据
        """
        try:
            import json
            import re
            
            # 尝试直接解析JSON
            try:
                # 查找JSON块
                json_match = re.search(r'```json\s*([\s\S]*?)\s*```', response_text)
                if json_match:
                    json_str = json_match.group(1)
                    data = json.loads(json_str)
                else:
                    # 尝试直接解析整个响应
                    data = json.loads(response_text)
                
                logger.info("成功解析Dify API返回的JSON数据")
            except json.JSONDecodeError:
                logger.warning("无法解析Dify API返回的JSON，尝试使用正则表达式提取")
                # 尝试从文本中提取信息
                data = self._extract_data_with_regex(response_text)
            
            # 生成任务ID
            task_id = str(uuid.uuid4())
            
            # 构造基本的任务数据
            task_data = {
                'task_id': task_id,
                'query_text': query_text
            }
            
            # 合并提取的数据
            if isinstance(data, dict):
                # 提取顶层字段到task_data
                for key in ['entry_url', 'category', 'description', 'priority', 'source_type']:
                    if key in data:
                        task_data[key] = data[key]
                
                # 提取子对象
                retrieval_conditions = data.get('retrieval_conditions', {})
                if retrieval_conditions:
                    retrieval_conditions['task_id'] = task_id
                
                scraping_strategies = data.get('scraping_strategies', {})
                if scraping_strategies:
                    scraping_strategies['task_id'] = task_id
                
                authorization_info = data.get('authorization_info', {})
                if authorization_info:
                    authorization_info['task_id'] = task_id
                
                return {
                    'task_data': task_data,
                    'retrieval_conditions': retrieval_conditions,
                    'scraping_strategies': scraping_strategies,
                    'authorization_info': authorization_info
                }
            else:
                logger.warning("解析后的数据不是字典类型")
                return {
                    'task_data': {
                        'task_id': task_id,
                        'query_text': query_text,
                        'description': "自动生成的任务"
                    }
                }
        except Exception as e:
            logger.error(f"解析Dify响应失败: {e}")
            # 返回一个基本的数据结构
            return {
                'task_data': {
                    'task_id': str(uuid.uuid4()),
                    'query_text': query_text,
                    'description': "解析失败，自动生成的任务"
                }
            }
    
    def _extract_data_with_regex(self, text):
        """
        使用正则表达式从文本中提取信息
        
        Args:
            text: 要提取信息的文本
        
        Returns:
            dict: 提取的信息
        """
        import re
        
        result = {}
        
        # 提取入口URL
        url_match = re.search(r'入口URL[^\n]*?[:：]\s*(https?://\S+)', text, re.IGNORECASE)
        if url_match:
            result['entry_url'] = url_match.group(1).strip()
        
        # 提取数据类别
        category_match = re.search(r'数据类别[^\n]*?[:：]\s*([^\n]+)', text, re.IGNORECASE)
        if category_match:
            result['category'] = category_match.group(1).strip()
        
        # 提取信源描述
        desc_match = re.search(r'信源描述[^\n]*?[:：]\s*([^\n]+)', text, re.IGNORECASE)
        if desc_match:
            result['description'] = desc_match.group(1).strip()
        
        # 提取优先级
        priority_match = re.search(r'优先级[^\n]*?[:：]\s*(高|中|低)', text, re.IGNORECASE)
        if priority_match:
            result['priority'] = priority_match.group(1).strip()
        
        # 提取信源类型
        type_match = re.search(r'信源类型[^\n]*?[:：]\s*([^\n]+)', text, re.IGNORECASE)
        if type_match:
            result['source_type'] = type_match.group(1).strip()
        
        # 提取检索条件部分
        retrieval_match = re.search(r'检索条件[^\n]*?[:：](.*?)(?=爬取策略|授权信息|$)', text, re.DOTALL)
        if retrieval_match:
            retrieval_text = retrieval_match.group(1).strip()
            retrieval_conditions = {}
            
            # 提取关键词
            keywords_match = re.search(r'关键词[^\n]*?[:：]\s*([^\n]+)', retrieval_text)
            if keywords_match:
                keywords_str = keywords_match.group(1).strip()
                retrieval_conditions['keywords'] = [k.strip() for k in keywords_str.split(',')]
            
            # 提取过滤条件
            filters_match = re.search(r'过滤条件[^\n]*?[:：]\s*([^\n]+)', retrieval_text)
            if filters_match:
                retrieval_conditions['filters'] = filters_match.group(1).strip()
            
            # 提取日期范围
            date_match = re.search(r'日期范围[^\n]*?[:：]\s*([^\n]+)', retrieval_text)
            if date_match:
                retrieval_conditions['date_range'] = date_match.group(1).strip()
            
            result['retrieval_conditions'] = retrieval_conditions
        
        # 提取爬取策略部分
        scraping_match = re.search(r'爬取策略[^\n]*?[:：](.*?)(?=授权信息|$)', text, re.DOTALL)
        if scraping_match:
            scraping_text = scraping_match.group(1).strip()
            scraping_strategies = {}
            
            # 提取爬取频率
            freq_match = re.search(r'爬取频率[^\n]*?[:：]\s*([^\n]+)', scraping_text)
            if freq_match:
                scraping_strategies['scraping_frequency'] = freq_match.group(1).strip()
            
            # 提取要提取的字段
            fields_match = re.search(r'要提取的字段[^\n]*?[:：]\s*([^\n]+)', scraping_text)
            if fields_match:
                fields_str = fields_match.group(1).strip()
                scraping_strategies['fields_to_extract'] = [f.strip() for f in fields_str.split(',')]
            
            result['scraping_strategies'] = scraping_strategies
        
        # 提取授权信息部分
        auth_match = re.search(r'授权信息[^\n]*?[:：](.*?)(?=$)', text, re.DOTALL)
        if auth_match:
            auth_text = auth_match.group(1).strip()
            authorization_info = {}
            
            # 提取API密钥
            api_key_match = re.search(r'API密钥[^\n]*?[:：]\s*([^\n]+)', auth_text)
            if api_key_match:
                authorization_info['api_key'] = api_key_match.group(1).strip()
            
            # 提取授权方式
            auth_method_match = re.search(r'授权方式[^\n]*?[:：]\s*([^\n]+)', auth_text)
            if auth_method_match:
                authorization_info['auth_method'] = auth_method_match.group(1).strip()
            
            result['authorization_info'] = authorization_info
        
        return result 